import random
import numpy as np
import math
import matplotlib.pyplot as plt

point_amount = 20

y = list(range(point_amount))
x = np.array(range(point_amount))
# x = np.array(np.zeros(point_amount))
for i in range(point_amount):
    y[i] = -i + random.uniform(-1,1)

y[int(point_amount/2)] = 20.0 #fake noise point

ob = np.array([x,y]).T

def normalize_angle(rad):
    '''
    归一化角度到 -PI ~ PI
    '''
    if abs(rad)<= math.pi:
        return rad
    else:
        ret = rad
        step = 0
        step = -2*math.pi if ret>0 else 2*math.pi
        while abs(ret) > math.pi:
            ret += step
        return ret

def reset_plot():
    plt.cla()
    plt.axis("equal")
    plt.grid(True)

def plot_data(m, color="or"):
    plt.plot(m[:, 0], m[:, 1], color, markersize=2)

def least_sqaure(ob):
    '''
    最小二乘法
    ===========

    - ob: 输入点必须是np.array([[x0, y0],[x1, y1], ... ,[x2, y2]])

    - return: (a, b, r, th) y = ax + b, r相关系数, th直线的角度
    '''
    x = ob[:,0]
    y = ob[:,1]
    meany = y.mean()
    meanx = x.mean()
    if abs(x.max()-meanx)<0.0001:
        print("Line has same x, return (None, meanx, 1.0, 90.0)")
        return (None, meanx, 1.0, 90.0)

    try:
        a = ((x*y).sum() - meany*x.sum()) / ((x**2).sum() - meanx*x.sum())
        b = (meany*(x**2).sum()-(x*y).sum()*meanx) / ((x**2).sum() - meanx*x.sum())
        r = ((x-meanx)*(y-meany)).sum() / math.sqrt(((x-meanx)**2).sum() * ((y-meany)**2).sum())
        th = math.degrees(normalize_angle(math.atan(a)))
    except BaseException as e:
        print(e)

    return (a,b,abs(r),th)

def best_line(points, min_relative=0.9):
    '''
    优化的最小二乘法, 过滤坏点, 直到相关性达到0.9的范围
    '''
    new_points = points
    line_args = [0, 0, 0]
    while line_args[2] < min_relative:
        line_args = least_sqaure(new_points)
        if len(new_points) < 10:
            # 少于10个点退出
            print('[WARN] two few points')
            break
        if line_args[0] is None:
            print('[Warn] Best Line is x=b')
            break
        func = lambda x: line_args[0] * x + line_args[1]
        err = abs(func(new_points[:,0]) - new_points[:,1])
        max_index = err.argmax()
        new_points = np.delete(new_points, max_index, 0)

    return line_args

line_arg = least_sqaure(ob)

line0 = np.array(
    [
        [x[0], 0],
        [x[len(x)-1],0]
    ]
)

line1 = np.array(
    [
        [x[0], 0],
        [x[len(x)-1],0]
    ]
)

def line_func(x:float, arg:list)->float:
    if arg[0] is None:
        return random.random()*100 #如果arg[0]是None，说明直线为x=b，随便返回一个y值
    else: 
        return arg[0] * x + arg[1]


line0[0,1] = line_func(line0[0,0], line_arg)
line0[1,1] = line_func(line0[1,0], line_arg)

line_arg = best_line(ob)
print("Best Line Args: "+str(line_arg))

line1[0,1] = line_func(line1[0,0], line_arg)
line1[1,1] = line_func(line1[1,0], line_arg)

reset_plot()
plot_data(ob)
plot_data(line0, "-b")
plot_data(line1, "-k")
plt.show()